"""
智能体注册中心 (Agent Registry)
支持动态发现、能力描述和智能体间直接通信
"""
from typing import Dict, Any, List, Optional, Set
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
import json
import uuid
from enum import Enum
from utils.logger import Logger
from middleware.redis_adapter import get_redis_adapter

logger = Logger.get_logger("agent_registry")


class AgentStatus(Enum):
    """智能体状态枚举"""
    ONLINE = "online"
    OFFLINE = "offline" 
    BUSY = "busy"
    ERROR = "error"
    MAINTENANCE = "maintenance"


@dataclass
class AgentCapability:
    """智能体能力描述"""
    name: str
    description: str
    input_schema: Dict[str, Any]  # JSON Schema
    output_schema: Dict[str, Any]  # JSON Schema
    performance_metrics: Dict[str, float] = None
    
    def __post_init__(self):
        if self.performance_metrics is None:
            self.performance_metrics = {
                "avg_response_time_ms": 0.0,
                "success_rate": 1.0,
                "throughput_qps": 0.0
            }


@dataclass
class AgentInfo:
    """智能体信息"""
    agent_id: str
    agent_type: str
    name: str
    version: str
    endpoint: str  # HTTP endpoint
    status: AgentStatus
    capabilities: List[AgentCapability]
    metadata: Dict[str, Any]
    registered_at: datetime
    last_heartbeat: datetime
    load_score: float = 0.0  # 负载评分，0-1
    reliability_score: float = 1.0  # 可靠性评分，0-1
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "agent_id": self.agent_id,
            "agent_type": self.agent_type,
            "name": self.name,
            "version": self.version,
            "endpoint": self.endpoint,
            "status": self.status.value,
            "capabilities": [
                {
                    "name": cap.name,
                    "description": cap.description,
                    "input_schema": cap.input_schema,
                    "output_schema": cap.output_schema,
                    "performance_metrics": cap.performance_metrics
                }
                for cap in self.capabilities
            ],
            "metadata": self.metadata,
            "registered_at": self.registered_at.isoformat(),
            "last_heartbeat": self.last_heartbeat.isoformat(),
            "load_score": self.load_score,
            "reliability_score": self.reliability_score
        }


class AgentRegistry:
    """智能体注册中心"""
    
    def __init__(self):
        self.redis_client = get_redis_adapter()
        self.registry_prefix = "agent_registry:"
        self.capability_index_prefix = "capability_index:"
        self.type_index_prefix = "type_index:"
        
    def register_agent(self, agent_info: AgentInfo) -> bool:
        """注册智能体"""
        try:
            # 存储智能体信息
            agent_key = f"{self.registry_prefix}{agent_info.agent_id}"
            agent_data = json.dumps(agent_info.to_dict())
            self.redis_client.setex(agent_key, 3600, agent_data)  # 1小时过期，需要心跳续期
            
            # 添加到类型索引
            type_key = f"{self.type_index_prefix}{agent_info.agent_type}"
            self.redis_client.sadd(type_key, agent_info.agent_id)
            self.redis_client.expire(type_key, 3600)
            
            # 添加到能力索引
            for capability in agent_info.capabilities:
                capability_key = f"{self.capability_index_prefix}{capability.name}"
                self.redis_client.sadd(capability_key, agent_info.agent_id)
                self.redis_client.expire(capability_key, 3600)
            
            logger.info(f"智能体注册成功: {agent_info.agent_id} ({agent_info.name})")
            return True
            
        except Exception as e:
            logger.error(f"智能体注册失败: {agent_info.agent_id}, error={str(e)}")
            return False
    
    def unregister_agent(self, agent_id: str) -> bool:
        """注销智能体"""
        try:
            # 获取智能体信息
            agent_info = self.get_agent_info(agent_id)
            if not agent_info:
                return False
            
            # 从类型索引中移除
            type_key = f"{self.type_index_prefix}{agent_info.agent_type}"
            self.redis_client.srem(type_key, agent_id)
            
            # 从能力索引中移除
            for capability in agent_info.capabilities:
                capability_key = f"{self.capability_index_prefix}{capability.name}"
                self.redis_client.srem(capability_key, agent_id)
            
            # 删除智能体信息
            agent_key = f"{self.registry_prefix}{agent_id}"
            self.redis_client.delete(agent_key)
            
            logger.info(f"智能体注销成功: {agent_id}")
            return True
            
        except Exception as e:
            logger.error(f"智能体注销失败: {agent_id}, error={str(e)}")
            return False
    
    def update_agent_status(self, agent_id: str, status: AgentStatus, 
                          load_score: Optional[float] = None,
                          reliability_score: Optional[float] = None) -> bool:
        """更新智能体状态"""
        try:
            agent_key = f"{self.registry_prefix}{agent_id}"
            agent_data = self.redis_client.get(agent_key)
            
            if not agent_data:
                logger.warning(f"智能体不存在: {agent_id}")
                return False
            
            agent_dict = json.loads(agent_data.decode('utf-8'))
            agent_dict["status"] = status.value
            agent_dict["last_heartbeat"] = datetime.utcnow().isoformat()
            
            if load_score is not None:
                agent_dict["load_score"] = load_score
            if reliability_score is not None:
                agent_dict["reliability_score"] = reliability_score
            
            self.redis_client.setex(agent_key, 3600, json.dumps(agent_dict))
            
            logger.debug(f"智能体状态更新: {agent_id} -> {status.value}")
            return True
            
        except Exception as e:
            logger.error(f"智能体状态更新失败: {agent_id}, error={str(e)}")
            return False
    
    def get_agent_info(self, agent_id: str) -> Optional[AgentInfo]:
        """获取智能体信息"""
        try:
            agent_key = f"{self.registry_prefix}{agent_id}"
            agent_data = self.redis_client.get(agent_key)
            
            if not agent_data:
                return None
            
            agent_dict = json.loads(agent_data.decode('utf-8'))
            
            capabilities = [
                AgentCapability(
                    name=cap["name"],
                    description=cap["description"],
                    input_schema=cap["input_schema"],
                    output_schema=cap["output_schema"],
                    performance_metrics=cap.get("performance_metrics", {})
                )
                for cap in agent_dict.get("capabilities", [])
            ]
            
            return AgentInfo(
                agent_id=agent_dict["agent_id"],
                agent_type=agent_dict["agent_type"],
                name=agent_dict["name"],
                version=agent_dict["version"],
                endpoint=agent_dict["endpoint"],
                status=AgentStatus(agent_dict["status"]),
                capabilities=capabilities,
                metadata=agent_dict.get("metadata", {}),
                registered_at=datetime.fromisoformat(agent_dict["registered_at"]),
                last_heartbeat=datetime.fromisoformat(agent_dict["last_heartbeat"]),
                load_score=agent_dict.get("load_score", 0.0),
                reliability_score=agent_dict.get("reliability_score", 1.0)
            )
            
        except Exception as e:
            logger.error(f"获取智能体信息失败: {agent_id}, error={str(e)}")
            return None
    
    def find_agents_by_capability(self, capability_name: str, 
                                  min_reliability: float = 0.8,
                                  max_load: float = 0.8) -> List[AgentInfo]:
        """根据能力查找智能体"""
        try:
            capability_key = f"{self.capability_index_prefix}{capability_name}"
            agent_ids = self.redis_client.smembers(capability_key)
            
            suitable_agents = []
            
            for agent_id_bytes in agent_ids:
                agent_id = agent_id_bytes.decode('utf-8')
                agent_info = self.get_agent_info(agent_id)
                
                if (agent_info and 
                    agent_info.status == AgentStatus.ONLINE and
                    agent_info.reliability_score >= min_reliability and
                    agent_info.load_score <= max_load):
                    suitable_agents.append(agent_info)
            
            # 按可靠性和负载排序
            suitable_agents.sort(key=lambda x: (x.reliability_score, -x.load_score), reverse=True)
            
            logger.info(f"能力匹配结果: {capability_name} -> {len(suitable_agents)} 个智能体")
            return suitable_agents
            
        except Exception as e:
            logger.error(f"能力查找失败: {capability_name}, error={str(e)}")
            return []
    
    def find_agents_by_type(self, agent_type: str) -> List[AgentInfo]:
        """根据类型查找智能体"""
        try:
            type_key = f"{self.type_index_prefix}{agent_type}"
            agent_ids = self.redis_client.smembers(type_key)
            
            agents = []
            for agent_id_bytes in agent_ids:
                agent_id = agent_id_bytes.decode('utf-8')
                agent_info = self.get_agent_info(agent_id)
                if agent_info:
                    agents.append(agent_info)
            
            return agents
            
        except Exception as e:
            logger.error(f"类型查找失败: {agent_type}, error={str(e)}")
            return []
    
    def get_system_overview(self) -> Dict[str, Any]:
        """获取系统概览"""
        try:
            # 获取所有智能体
            pattern = f"{self.registry_prefix}*"
            keys = self.redis_client.keys(pattern)
            
            total_agents = len(keys)
            status_counts = {status.value: 0 for status in AgentStatus}
            type_counts = {}
            capability_counts = {}
            
            for key in keys:
                agent_data = self.redis_client.get(key)
                if agent_data:
                    agent_dict = json.loads(agent_data.decode('utf-8'))
                    
                    # 统计状态
                    status = agent_dict.get("status", "unknown")
                    if status in status_counts:
                        status_counts[status] += 1
                    
                    # 统计类型
                    agent_type = agent_dict.get("agent_type", "unknown")
                    type_counts[agent_type] = type_counts.get(agent_type, 0) + 1
                    
                    # 统计能力
                    for cap in agent_dict.get("capabilities", []):
                        cap_name = cap.get("name", "unknown")
                        capability_counts[cap_name] = capability_counts.get(cap_name, 0) + 1
            
            return {
                "total_agents": total_agents,
                "status_distribution": status_counts,
                "type_distribution": type_counts,
                "capability_distribution": capability_counts,
                "timestamp": datetime.utcnow().isoformat()
            }
            
        except Exception as e:
            logger.error(f"获取系统概览失败: error={str(e)}")
            return {}
    
    def cleanup_offline_agents(self, offline_threshold: int = 300) -> int:
        """清理离线智能体"""
        try:
            pattern = f"{self.registry_prefix}*"
            keys = self.redis_client.keys(pattern)
            
            cleaned_count = 0
            current_time = datetime.utcnow()
            
            for key in keys:
                agent_data = self.redis_client.get(key)
                if agent_data:
                    agent_dict = json.loads(agent_data.decode('utf-8'))
                    
                    last_heartbeat = datetime.fromisoformat(agent_dict["last_heartbeat"])
                    offline_duration = (current_time - last_heartbeat).total_seconds()
                    
                    if offline_duration > offline_threshold:
                        agent_id = agent_dict["agent_id"]
                        self.unregister_agent(agent_id)
                        cleaned_count += 1
            
            logger.info(f"清理离线智能体: {cleaned_count} 个")
            return cleaned_count
            
        except Exception as e:
            logger.error(f"清理离线智能体失败: error={str(e)}")
            return 0


# 全局注册中心实例
agent_registry = AgentRegistry()


# 使用示例
def example_usage():
    """注册中心使用示例"""
    
    # 创建智能体能力
    text_qa_capability = AgentCapability(
        name="text_question_answering",
        description="文本问答能力，支持自然语言查询",
        input_schema={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "用户查询"},
                "context": {"type": "string", "description": "可选上下文"}
            },
            "required": ["query"]
        },
        output_schema={
            "type": "object",
            "properties": {
                "answer": {"type": "string", "description": "答案"},
                "confidence": {"type": "number", "description": "置信度"},
                "sources": {"type": "array", "description": "参考源"}
            }
        }
    )
    
    # 创建智能体信息
    agent_info = AgentInfo(
        agent_id="text_qa_agent_001",
        agent_type="text_qa",
        name="文本问答智能体",
        version="1.0.0",
        endpoint="http://localhost:8001",
        status=AgentStatus.ONLINE,
        capabilities=[text_qa_capability],
        metadata={"model": "gpt-3.5-turbo", "max_tokens": 4096},
        registered_at=datetime.utcnow(),
        last_heartbeat=datetime.utcnow()
    )
    
    # 注册智能体
    success = agent_registry.register_agent(agent_info)
    print(f"注册结果: {success}")
    
    # 查找具有特定能力的智能体
    capable_agents = agent_registry.find_agents_by_capability("text_question_answering")
    print(f"找到 {len(capable_agents)} 个具有文本问答能力的智能体")
    
    # 获取系统概览
    overview = agent_registry.get_system_overview()
    print(f"系统概览: {overview}")


if __name__ == "__main__":
    example_usage()